#!/usr/bin/python
# -*- coding: UTF-8 -*-
import base64
import json
import os
import time
import sys

import win32gui

import cv2
import win32con
from PIL import Image

import requests
from selenium import webdriver
from pywinauto import mouse, keyboard
from utils import *
from selenium.common.exceptions import TimeoutException, NoAlertPresentException
from selenium.webdriver.support.ui import Select
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support.expected_conditions import alert_is_present

catalogMap = {
    "工程科技Ⅰ辑": "Bfirst"
}


class Crawler:
    def __init__(self, user_name, password, catalog_name, topic, page_year):
        self.user_name = user_name
        self.password = password
        self.catalog_name = catalog_name
        self.topic = topic
        self.page_year = page_year
        self.save_file_name = "./result/{}-{}.txt".format(topic, page_year)
        self.driver = None

    def select_paper(self):
        # 清除已选
        self.driver.find_element_by_xpath('//*[@id="XuekeNavi_Div"]/div[1]/input[1]').click()
        time.sleep(2)
        print(catalogMap[self.catalog_name])
        self.driver.find_element_by_id(catalogMap[self.catalog_name]).click()
        time.sleep(2)
        self.driver.find_element_by_name(self.topic).click()
        time.sleep(2)

        # 检索论文
        self.driver.find_element_by_id('btnSearch').click()
        time.sleep(5)
        # self.driver.find_element_by_link_text(u"发表年度").click()
        self.driver.find_element_by_xpath('//*[@id="alink1"]').click()
        time.sleep(5)

        self.driver.find_element_by_link_text(self.page_year).click()
        time.sleep(5)

    # def login_cnki(self):
    #     time.sleep(1)
    #     driver.find_element_by_id("Ecp_top_login_show").click()
    #     driver.find_element_by_id("Ecp_TextBoxUserName").send_keys(user_name)
    #     driver.find_element_by_id("Ecp_TextBoxPwd").send_keys(password)
    #     driver.find_element_by_id("Ecp_Button1").click()

    def login_haust(self):
        time.sleep(1)
        self.driver.find_element_by_id("username").send_keys(self.user_name)
        self.driver.find_element_by_id("password").send_keys(self.password)
        self.driver.find_element_by_xpath('//*[@id="fm1"]/table/tbody/tr[4]/td/input').click()
        time.sleep(2)

        self.driver.find_element_by_xpath('//*[@id="collapse-panel"]/div/div/div[1]/div/div/input').send_keys("https")
        time.sleep(1)
        self.driver.find_element_by_id("quick-access-input").send_keys("https://kns.cnki.net/kns/brief/result.aspx")
        time.sleep(1)
        self.driver.find_element_by_xpath('//*[@id="go"]').click()
        time.sleep(2)
        self.driver.close()
        time.sleep(1)
        self.driver.switch_to.window(self.driver.window_handles[0])

    def login_haust_new(self):
        self.driver.find_element_by_xpath('/html/body/div[1]/div[1]/div/div/ul/li[1]/a').click()
        time.sleep(2)
        self.driver.close()
        time.sleep(2)

        self.driver.switch_to.window(self.driver.window_handles[0])
        self.driver.find_element_by_xpath('//*[@id="highSearch"]').click()
        time.sleep(2)
        self.driver.close()
        time.sleep(2)
        self.driver.switch_to.window(self.driver.window_handles[0])
        time.sleep(2)

    def base64_api(self, uname, pwd, img, typeid):
        with open(img, 'rb') as f:
            base64_data = base64.b64encode(f.read())
            b64 = base64_data.decode()
        data = {"username": uname, "password": pwd, "typeid": typeid, "image": b64}
        result = json.loads(requests.post("http://api.ttshitu.com/predict", json=data).text)
        if result['success']:
            return result["data"]["result"]
        else:
            return result["message"]
        return ""

    def get_verification_code(self, img_element, screen_image, code_image, iframe_x, iframe_y):
        location = img_element.location
        size = img_element.size
        x1 = int(location['x']) + iframe_x
        y1 = int(location['y']) + iframe_y
        x2 = int(location['x'] + size['width']) + iframe_x
        y2 = int(location['y'] + size['height']) + iframe_y
        rangle = (x1, y1, x2, y2)  # 写成我们需要截取的位置坐标

        i = Image.open(screen_image)  # 打开截图
        frame4 = i.crop(rangle)
        frame4.save(code_image)

        # img = cv2.imread(screen_image)  # 读取目录下的jpg图像
        # cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0))
        # cv2.imshow('real_image', img)
        # k = cv2.waitKey(0)  # waitKey代表读取键盘的输入，0代表一直等待
        # if k == 27:  # 键盘上Esc键的键值
        #     cv2.destroyAllWindows()

        # api识别部分
        result = self.base64_api('ZebinFang', '83731580', code_image, 3)
        # 输出识别出的验证码
        print("验证码识别：", result)
        return result

    # def download_pdf(self):
    #     self.driver.find_element_by_id("pdfDown").click()
    #     time.sleep(2)
    #     is_downloaded = True
    #
    #     if len(self.driver.window_handles) == 3:
    #         self.driver.switch_to.window(self.driver.window_handles[2])
    #
    #         try:
    #             alert1 = self.driver.switch_to.alert
    #             alert1.accept()
    #             element_pic = img_2_location(0.7)
    #             if element_pic is not None:
    #                 print("match image success")
    #                 confirm_button = [element_pic['rectangle'][3][0], element_pic['rectangle'][3][1]]
    #                 mouse.click('left', confirm_button)
    #                 time.sleep(1)
    #                 is_downloaded = False
    #             elif len(self.driver.find_elements_by_link_text(u'检索服务器暂时没有提供服务,请您重新登录到其它服务器！')) == 1:
    #                 print("close all window, login again")
    #                 self.driver.quit()
    #                 is_downloaded = False
    #                 return False, is_downloaded
    #             else:
    #                 self.driver.find_element_by_xpath('//*[@id="vImg"]')
    #                 screen_image = './resource/image/screen.png'
    #                 self.driver.save_screenshot(screen_image)
    #                 img_element = self.driver.find_element_by_id("vImg")
    #
    #                 code_image = "./resource/image/code.png"
    #                 result = self.get_verification_code(img_element, screen_image, code_image, 0, 0)
    #
    #                 # 将验证码填入
    #                 self.driver.find_element_by_id('vcode').send_keys(result)
    #                 time.sleep(1)
    #
    #                 # 点击验证码的提交，然后开始下载
    #                 self.driver.find_element_by_xpath('/html/body/div/form/dl/dd/button').click()
    #                 time.sleep(3)
    #                 self.driver.close()
    #                 time.sleep(1)
    #                 self.driver.switch_to.window(self.driver.window_handles[1])
    #             time.sleep(1)
    #         except:
    #             print("can not handle image")
    #             is_downloaded = False
    #             self.driver.switch_to.window(self.driver.window_handles[1])
    #             time.sleep(2)
    #     else:
    #         time.sleep(5)
    #
    #     self.driver.close()
    #     time.sleep(1)
    #     self.driver.switch_to.window(self.driver.window_handles[0])
    #     time.sleep(1)
    #     self.driver.switch_to.frame("iframeResult")
    #     return True, is_downloaded

    def download_pdf(self):
        self.driver.find_element_by_id("pdfDown").click()
        time.sleep(8)
        is_downloaded = True

        if len(self.driver.window_handles) == 3:
            self.driver.switch_to.window(self.driver.window_handles[2])
            try:
                print("exist alert")
                # alert1 = self.driver.switch_to.alert
                # alert1.accept()
            except NoAlertPresentException:
                print("no exist alert")
                self.driver.find_element_by_xpath('//*[@id="vImg"]')
                screen_image = './resource/image/screen.png'
                self.driver.save_screenshot(screen_image)
                img_element = self.driver.find_element_by_id("vImg")

                code_image = "./resource/image/code.png"
                result = self.get_verification_code(img_element, screen_image, code_image, 0, 0)

                # 将验证码填入
                self.driver.find_element_by_id('vcode').send_keys(result)
                time.sleep(1)

                # 点击验证码的提交，然后开始下载
                self.driver.find_element_by_xpath('/html/body/div/form/dl/dd/button').click()
                time.sleep(3)
                self.driver.close()
                time.sleep(1)
                self.driver.switch_to.window(self.driver.window_handles[1])
                time.sleep(1)
            except:
                print("can not handle image")
                is_downloaded = False
                self.driver.switch_to.window(self.driver.window_handles[1])
                time.sleep(2)
        else:
            time.sleep(5)

        self.driver.close()
        time.sleep(1)
        self.driver.switch_to.window(self.driver.window_handles[0])
        time.sleep(1)
        self.driver.switch_to.frame("iframeResult")
        return True, is_downloaded

    def download_page_paper(self, page_num):
        index = 0
        current_num = 1

        paper_table = self.driver.find_element_by_class_name("GridTableContent")
        paper_rows = paper_table.find_elements_by_tag_name('tr')
        paper_name_list = ["current page: {}\n".format(page_num)]

        try:
            for row in paper_rows:
                if index == 0:
                    index += 1
                    continue

                if index < current_num:
                    print("current page: " + str(page_num) + ", skip row: " + str(index))
                    index += 1
                    continue

                if len(row.find_elements_by_class_name("briefDl_D")) == 0 \
                        and len(row.find_elements_by_class_name("briefDl_Y")) == 0:
                    print("current page: " + str(page_num) + ", row: " + str(index) + ", no download button")
                    index += 1
                    time.sleep(0.5)
                    continue

                if len(row.find_elements_by_class_name("briefDl_N")) == 1:
                    print("current page: " + str(page_num) + ", row: " + str(index) + ", the product is out of range")
                    index += 1
                    time.sleep(0.5)
                    continue

                print("current page: " + str(page_num) + ", row: " + str(index))

                row.find_element_by_class_name("fz14").click()
                time.sleep(3)
                self.driver.switch_to.window(self.driver.window_handles[1])

                re_login, is_downloaded = self.download_pdf()
                if is_downloaded:
                    paper_name = row.find_element_by_class_name("fz14").text
                    # journal_name = row.find_element_by_class_name("cjfdyxyz").find_element_by_tag_name("a").text
                    # paper_name_list.append("{},{}\n".format(paper_name, journal_name))
                    paper_name_list.append("{}\n".format(paper_name))

                if not re_login:
                    return False

                index += 1
                time.sleep(3)
        finally:
            # TODO 论文提取
            with open(self.save_file_name, "a+") as f:
                f.writelines(paper_name_list)
        return True

    def handle_page_each_verification_code(self, iframe_x, iframe_y):
        self.driver.find_element_by_xpath('//*[@id="CheckCodeImg"]')
        screen_image = './resource/image/page_screen.png'
        self.driver.save_screenshot(screen_image)
        img_element = self.driver.find_element_by_id('CheckCodeImg')

        code_image = "./resource/image/page_code.png"
        result = self.get_verification_code(img_element, screen_image, code_image, iframe_x, iframe_y)

        # 将验证码填入
        self.driver.find_element_by_id('CheckCode').send_keys(result)
        time.sleep(1)

        # 点击验证码的提交，然后开始下载
        self.driver.find_element_by_xpath('/html/body/p[1]/input[2]').click()
        time.sleep(3)

    def handle_page_verification_code(self, iframe_x, iframe_y):
        if len(self.driver.find_elements_by_link_text('下一页')) == 0:
            self.handle_page_each_verification_code(iframe_x, iframe_y)
            time.sleep(2)

            if len(self.driver.find_elements_by_link_text('下一页')) == 0:
                # handle_page_each_verification_code(iframe_x, iframe_y)
                # time.sleep(2)
                print("filling verification code manually")

    def download_paper(self, start_page):
        iframe_location =self.driver.find_element_by_id("iframeResult").location
        iframe_x = int(iframe_location['x'])
        iframe_y = int(iframe_location['y'])

        self.driver.switch_to.frame("iframeResult")

        self.driver.find_element_by_xpath('//*[@id="id_grid_display_num"]/a[3]').click()
        time.sleep(1)

        self.handle_page_verification_code(iframe_x, iframe_y)

        page_info = self.driver.find_element_by_class_name("countPageMark").text
        page_list = page_info.split("/")
        total_page = int(page_list[1])

        current_page = 1

        if start_page > 9:
            page_table = self.driver.find_element_by_class_name("TitleLeftCell")
            print("first jump to page 9")
            page_table.find_element_by_link_text("9").click()
            current_page = 9

        while True:
            if current_page > total_page:
                return True, total_page

            # 处理翻页验证码
            self.handle_page_verification_code(iframe_x, iframe_y)

            # 下载当前页paper
            if current_page >= start_page:
                if not self.download_page_paper(current_page):
                    return False, current_page+1
                current_page + 1

            elif current_page >= 5:  # 一次翻多页
                if current_page + 4 <= start_page:
                    current_page += 4
                else:
                    current_page = start_page

                if current_page > total_page:
                    return True, total_page

                print("jump to page " + str(current_page))
                page_table = self.driver.find_element_by_class_name("TitleLeftCell")
                page_table.find_element_by_link_text(str(current_page)).click()
                time.sleep(3)
                continue

            self.driver.find_element_by_class_name("TitleLeftCell").find_element_by_link_text('下一页').click()
            time.sleep(3)
            current_page += 1

        # 返回到原界面
        self.driver.switch_to.default_content()

    def handle_paper(self, first_start_page):
        # self.login_haust()
        # self.select_paper()
        # is_finished, page_num = self.download_paper(first_start_page)
        #
        # if not is_finished:
        #     print("login again, from page: %d" % page_num)
        #     self.init_driver()
        #     self.handle_paper(page_num)
        # else:
        #     print("total page %d has downloaded" % page_num)
        self.login_haust_new()
        self.select_paper()
        self.download_paper(first_start_page)

    def init_driver(self):
        opts = webdriver.ChromeOptions()
        # opts.add_argument('--headless')
        opts.binary_location = "C:\Program Files\Google\Chrome\Application\chrome.exe"
        opts.add_experimental_option("excludeSwitches", ['enable-automation'])
        driver = webdriver.Chrome(executable_path="./resource/chromedriver.exe", options=opts)
        driver.maximize_window()
        # driver.get('http://wvpn.haust.edu.cn/login')
        driver.get('https://www.cnki.net/')
        self.driver = driver
